"""
Based on https://github.com/asanakoy/kaggle_carvana_segmentation
"""
import torch
import torch.utils.data as data
from torch.autograd import Variable as V

# import cv2
import numpy as np
import os

import pandas as pd

MODEL_MIN = 200
MODEL_MAX = 2500

DATA_MIN = 0
DATA_MAX = 1.4


def data_loader(id, root, page):

    model = np.array(pd.read_csv(os.path.join(root, '{}.bmdl').format(id), sep='\\s+', header=None), dtype=np.float32)
    model = model[:, 0].T
    model = np.expand_dims(model, axis=0)
    model = 2*(model - MODEL_MIN) / (MODEL_MAX - MODEL_MIN) - 1

    data = np.array(pd.read_csv(os.path.join(root, '{}.dat').format(id), sep='\\s+', header=None), dtype=np.float32)
    data = data[:, 1].T
    data = np.expand_dims(data, axis=0)
    data = 2*(data - DATA_MIN) / (DATA_MAX - DATA_MIN) - 1

    return model, data

def data_loader_20(id, root, page):

    model = np.array(pd.read_csv(os.path.join(root, '{}.bmdl').format(id), sep='\\s+', header=None), dtype=np.float32)
    model = model[:, :].T
    #model = np.expand_dims(model, axis=0)
    model = 2*(model - MODEL_MIN) / (MODEL_MAX - MODEL_MIN) - 1

    data = np.array(pd.read_csv(os.path.join(root, '{}.dat').format(id), sep='\\s+', header=None), dtype=np.float32)
    data = data[:, 1].T
    data = np.expand_dims(data, axis=0)
    data = 2*(data - DATA_MIN) / (DATA_MAX - DATA_MIN) - 1

    return model, data


class Dataset(data.Dataset):

    def __init__(self, ids, root, color, eval=False):
        self.ids = ids
        self.loader = data_loader
        self.root = root
        self.color = color
        self.eval = eval

    def __getitem__(self, index):
        id = self.ids[index]
        model, data = self.loader(id, self.root, 0)
        model = torch.Tensor(model)
        data = torch.Tensor(data)
        return id, model, data

    def __len__(self):
        return len(list(self.ids))

class Dataset_20(data.Dataset):

    def __init__(self, ids, root, color, eval=False):
        self.ids = ids
        self.loader = data_loader_20
        self.root = root
        self.color = color
        self.eval = eval

    def __getitem__(self, index):
        id = self.ids[index]
        model, data = self.loader(id, self.root, 0)
        model = torch.Tensor(model)
        data = torch.Tensor(data)
        return id, model, data

    def __len__(self):
        return len(list(self.ids))